package com.demo.demo3.utils;

import lombok.AllArgsConstructor;
import lombok.Data;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

public class TreeUtil {

    @Data
    @AllArgsConstructor
    public static class TreeVO<T, K> {
        private K parentKey;
        private K key;
        private String code;
        private T data;
        private List<TreeVO<T, K>> children;
    }

    /**
     * 列表转树
     */
    public static <T, K> List<TreeVO<T, K>> listToTree(
            List<T> list,
            Function<T, K> getKey,
            Function<K, K> getParentKey
    ) {
        // 构建节点列表
        List<TreeVO<T, K>> nodes = list
                .stream()
                .map(item -> {
                    K key = getKey.apply(item);
                    return new TreeVO<>(
                            getParentKey.apply(key),
                            key,
                            null,
                            item,
                            null
                    );
                })
                .collect(Collectors.toList());
        // 获取所有节点Key
        Set<K> keySet = nodes.stream().map(TreeVO::getKey).collect(Collectors.toSet());
        // 递归构建不存在的节点
        for (TreeVO<T, K> item : new ArrayList<>(nodes)) {
            createNotExistNode(nodes, keySet, item.getKey(), getParentKey);
        }
        // 构建每个节点的子节点列表
        Map<K, List<TreeVO<T, K>>> childrenMap = new HashMap<>();
        for (TreeVO<T, K> item : nodes) {
            K parentKey = item.parentKey;
            if (!childrenMap.containsKey(parentKey)) {
                childrenMap.put(parentKey, new ArrayList<>());
            }
            childrenMap.get(parentKey).add(item);
        }
        // 补充每个节点的子节点列表
        for (TreeVO<T, K> item : nodes) {
            K key = item.key;
            item.setChildren(childrenMap.get(key));
        }
        // 获取根节点
        List<TreeVO<T, K>> rootTrees = childrenMap.get(null);
        // 补充根节点序号
        createCode(rootTrees, new ArrayList<>());
        // 返回根节点
        return rootTrees;
    }

    /**
     * 补充缺失节点
     */
    private static <T, K> void createNotExistNode(List<TreeVO<T, K>> nodes, Set<K> keySet, K key, Function<K, K> getParentKey) {
        K parentKey = getParentKey.apply(key);
        if (parentKey == null) {
            return;
        }
        if (!keySet.contains(parentKey)) {
            keySet.add(parentKey);
            nodes.add(new TreeVO<>(
                    getParentKey.apply(parentKey),
                    parentKey,
                    null,
                    null,
                    null
            ));
        }
        createNotExistNode(nodes, keySet, parentKey, getParentKey);
    }

    /**
     * 补充树层级编号
     */
    private static <T, K> void createCode(List<TreeVO<T, K>> list, List<String> parentCodes) {
        if (list == null || list.isEmpty()) {
            return;
        }
        for (int i = 0; i < list.size(); i++) {
            TreeVO<T, K> item = list.get(i);
            List<String> codes = new ArrayList<>(parentCodes);
            codes.add(i + 1 + "");
            String code = String.join(".", codes);
            item.setCode(code);
            createCode(item.getChildren(), codes);
        }
    }

}
